package com.xxx.interceptor;

import java.lang.reflect.Method;
import java.util.concurrent.TimeUnit;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.HttpStatus;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.ModelAndView;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;

import com.xxx.annotation.RateLimiterAnno;
import com.xxx.annotation.RateLimiterAnno.TYPE;
import com.xxx.ratelimiter.GuavaRateLimiter;
import com.xxx.ratelimiter.GuavaRateLimiterFactory;

/**
 * @RateLimiter 限流注解拦截器
 * (需要springmvc.xml配置拦截生效)
 */
public class RateLimitInterceptor extends HandlerInterceptorAdapter {

	private static final Logger logger = LoggerFactory.getLogger(RateLimitInterceptor.class);

	private static final String[] IP_HEADER_KYES = { 
			"X-Forwarded-For", 
			"X-Real-IP", 
			"Proxy-Client-IP",
			"WL-Proxy-Client-IP", 
			"HTTP_CLIENT_IP", 
			"HTTP_X_FORWARDED_FOR" };

	private static final String USER_TOKEN_KEY = "USER_TOKEN";

	@Autowired
	private GuavaRateLimiterFactory rateLimiterFactory;

	@Override
	public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler)
			throws Exception {
		logger.info("enter RateLimitInterceptor preHandle()...");
		if (!(handler instanceof HandlerMethod)) {
			return true;
		}
		
		// 是否被限流（true：不限流，false：限流）
		boolean result = true;
		HandlerMethod handlerMethod = (HandlerMethod) handler;
		Method method = handlerMethod.getMethod();
		// 判断方法上是否有限流器注解：@RateLimiterAnno
		if (method.isAnnotationPresent(RateLimiterAnno.class)) {
			// 从注解上获取限流相关属性值（也可以通过其他方式获取限流配置参数，例如配置文件、数据库等）
			RateLimiterAnno rateLimiterAnnotation = method.getAnnotation(RateLimiterAnno.class);
			String apiKey = rateLimiterAnnotation.apiKey(); 		// 限流的api key
			TYPE type = rateLimiterAnnotation.type();  				// 限流策略类型
			int permits = rateLimiterAnnotation.permits();  		// 限流次数
			TimeUnit timeUnit = rateLimiterAnnotation.timeUnit(); 	// 时间单位
			
			if ("".equals(apiKey)) { // 如果api key为空，默认使用RequestURI作为key
				apiKey = request.getRequestURI();
			}

			// 根据不同的限流策略，执行对应的限流措施
			String rateLimiterKey = this.getRateLimiterKey(request, apiKey, type); // 限流key
			
			// 限流处理
			result = this.checkLimit(rateLimiterKey, timeUnit, permits, response);
		}
		
		return result;
	}

	/**
	 * 校验是否超过限流阀值
	 * @param rateLimiterKey
	 * @param timeUnit
	 * @param permits
	 * @param response
	 * @return true：拿到了token，未被限流；false：未拿到token，被限流
	 * @throws Exception
	 */
	private boolean checkLimit(String rateLimiterKey, TimeUnit timeUnit, int permits,
			HttpServletResponse response) throws Exception {
		GuavaRateLimiter rateLimiter = rateLimiterFactory.get(rateLimiterKey, timeUnit, permits);
		if (!rateLimiter.acquire(rateLimiterKey)) { // 调用acquire()方法判断是否被限流
			logger.error("rateLimiterKey:{} Access denied.", rateLimiterKey);
			response.setStatus(HttpStatus.FORBIDDEN.value()); // 被限流，返回禁止请求
			response.getWriter().print("Access denied because of exceeding access rate");
			return false;
		}
		return true;
	}

	/**
	 * 根据不同的限流策略，执行对应的限流措施
	 * @param request
	 * @param apiKey api key
	 * @param type 限流策略类型
	 * @return 限流key
	 */
	private String getRateLimiterKey(HttpServletRequest request, String apiKey, TYPE type) {
		String rateLimiterKey = "";
		if (RateLimiterAnno.TYPE.ALL == type) { // 整体限流
			rateLimiterKey = apiKey;
			
		} else if (RateLimiterAnno.TYPE.IP == type) { // 按IP限流
			String ip = getIP(request); // 获取请求的客户端IP
			if (ip != null) {
				rateLimiterKey = apiKey + ":" + ip;
			}
			
		} else if (RateLimiterAnno.TYPE.USER == type) {  // 按单个用户限流
			String userToken = getUserToken(request); // 获取登录用户key
			if (userToken != null) {
				rateLimiterKey = apiKey + ":" + userToken;
			}
		}
		return rateLimiterKey;
	}

	// 获取客户端IP
	private String getIP(HttpServletRequest request) {
		for (String ipHeaderKey : IP_HEADER_KYES) {
			String ip = request.getHeader(ipHeaderKey);
			if (ip != null && ip.length() != 0 && (!"unknown".equalsIgnoreCase(ip))) {
				return ip;
			} else {
				return request.getRemoteHost();
			}
		}
		return null;
	}

	// 获取登录用户token（或id）
	private String getUserToken(HttpServletRequest request) {
		String userToken = request.getHeader(USER_TOKEN_KEY);
		return userToken;
	}

	@Override
	public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler,
			ModelAndView modelAndView) throws Exception {
	}

	@Override
	public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler,
			Exception ex) throws Exception {
	}

}
